use std::collections::HashMap;
use std::sync::{ Arc, Mutex };
use lazy_static::lazy_static;

// Define your hosts
lazy_static! {
    static ref EASTMONEY_SSE_HOSTS: Vec<String> = {
        let mut hosts = Vec::new();
        for i in 1..500 {
            hosts.push(format!("{}.push2.eastmoney.com", i));
        }
        hosts
    };

    static ref EASTMONEY_SSE_DELAY_HOSTS: Vec<String> = {
        let mut hosts = Vec::new();
        for i in 1..500 {
            hosts.push(format!("{}.push2delay.eastmoney.com", i));
        }
        hosts
    };
}

// Create instances of `EvenlyDistributedHosts`
lazy_static! {
    pub static ref TICK3_HOST_MANAGER: EvenlyDistributedHosts = EvenlyDistributedHosts::new(
        EASTMONEY_SSE_HOSTS.clone()
    );

    pub static ref M001KLINES_HOST_MANAGER: EvenlyDistributedHosts = EvenlyDistributedHosts::new(
        EASTMONEY_SSE_HOSTS.clone()
    );

    pub static ref DELAY_HOST_MANAGER: EvenlyDistributedHosts = EvenlyDistributedHosts::new(
        EASTMONEY_SSE_DELAY_HOSTS.clone()
    );
}

/// A structure to manage the distribution of security IDs across hosts.
pub struct EvenlyDistributedHosts {
    security_id_host_map: Arc<Mutex<HashMap<String, String>>>,
    host_security_ids: Arc<Mutex<HashMap<String, Vec<String>>>>,
    hosts: Vec<String>,
}

impl EvenlyDistributedHosts {
    /// Creates a new instance of `EvenlyDistributedHosts`.
    fn new(hosts: Vec<String>) -> Self {
        let security_id_host_map = Arc::new(Mutex::new(HashMap::new()));
        let host_security_ids = Arc::new(
            Mutex::new(
                hosts
                    .iter()
                    .map(|host| (host.clone(), Vec::new()))
                    .collect()
            )
        );

        EvenlyDistributedHosts {
            security_id_host_map,
            host_security_ids,
            hosts,
        }
    }

    /// Distributes a security ID evenly across the available hosts.
    /// Returns the host to which the given security_id should connect.
    pub fn add_to_hosts_evenly(&self, security_id: &str) -> String {
        let mut host_security_ids = self.host_security_ids.lock().unwrap();
        let mut security_id_host_map = self.security_id_host_map.lock().unwrap();

        let min_host = self.hosts
            .iter()
            .min_by_key(|host| host_security_ids.get(&host.to_string()).unwrap().len())
            .unwrap();

        let host = min_host.clone();
        host_security_ids.get_mut(&host).unwrap().push(security_id.to_string());
        security_id_host_map.insert(security_id.to_string(), host.clone());

        host
    }

    /// Retrieves the host for a given security_id.
    pub fn get_host_for_security_id(&self, security_id: &str) -> Option<String> {
        let security_id_host_map = self.security_id_host_map.lock().unwrap();
        security_id_host_map.get(security_id).cloned()
    }

    /// Retrieves all security_ids for a given host.
    fn get_security_ids_for_host(&self, host: &str) -> Vec<String> {
        let host_security_ids = self.host_security_ids.lock().unwrap();
        host_security_ids.get(host).cloned().unwrap_or_default()
    }
}

#[cfg(test)]
mod test {
    use std::collections::HashSet;

    use crate::todo::evenly_hosts::{ DELAY_HOST_MANAGER, TICK3_HOST_MANAGER };

    use super::EvenlyDistributedHosts;

    // Function to generate a set of unique security IDs for testing.
    fn generate_security_ids(count: usize) -> Vec<String> {
        (0..count).map(|i| format!("security_{}", i)).collect()
    }

    // Function to test the distribution of security IDs and retrieval of information.
    fn test_distribution(host_manager: &EvenlyDistributedHosts, security_ids_count: usize) {
        let security_ids = generate_security_ids(security_ids_count);

        // Distribute the security IDs evenly across the hosts.
        for security_id in &security_ids {
            host_manager.add_to_hosts_evenly(security_id);
        }

        // Check that each security ID has been assigned to a host.
        let mut assigned_hosts = HashSet::new();
        for security_id in &security_ids {
            if let Some(host) = host_manager.get_host_for_security_id(security_id) {
                assigned_hosts.insert(host);
            }
        }

        println!("Number of distinct hosts assigned: {}", assigned_hosts.len());

        // Verify that all hosts have been assigned security IDs.
        for host in &host_manager.hosts {
            let ids = host_manager.get_security_ids_for_host(host);
            assert!(!ids.is_empty(), "Host {} has no security IDs.", host);
        }
    }

    #[tokio::test]
    async fn test_main() {
        // Test the distribution with SSE hosts.
        test_distribution(&TICK3_HOST_MANAGER, 1000);

        // Test the distribution with delay hosts.
        test_distribution(&DELAY_HOST_MANAGER, 1000);
    }
}
